{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Concepts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`metadsl` inserts a layer between calling a function and computing its result, so that we can build up a bunch of calls, transform them, and then execute them all at once."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Expression\n",
    "\n",
    "We start with a building block `Expression`s:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import dataclasses\n",
    "import metadsl\n",
    "\n",
    "class MyObject(metadsl.Expression):\n",
    "    @metadsl.expression\n",
    "    def do_things(self) -> MyObject:\n",
    "        ...\n",
    "\n",
    "    @metadsl.expression\n",
    "    def __add__(self, other: MyObject) -> MyObject:\n",
    "        ...\n",
    "\n",
    "        \n",
    "@metadsl.expression\n",
    "def create_object(x: int) -> MyObject:\n",
    "    ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "create_object(123)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o = create_object(123)\n",
    "o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "myobject_0 = create_object(123)\n",
       "myobject_0.do_things().__add__(myobject_0)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o.do_things() + o"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is useful to keep in mind the strict typing constraints here, not all of which can be faithfully checked by MyPy:\n",
    "\n",
    "1. The arguments in a `Expression` should fit the signature of the `_function`.\n",
    "2. The return type of the `_function` should correspond to the type of the `Expression`.\n",
    "\n",
    "Let's check the first two of these for `o`. We see that the expression's function is `create_object`, which takes in an `int` and returns a `MyObject`. What is `metadsl.E`? It is a type alias for `Union[T, metadsl.LiteralExpression[T]]`. This represents anything that could be a python literal, or a leaf of the expression tree. The argument here is an `int`, which is compatible with the argument hint. And the instances holding it is of type `MyObject`, which is its return type."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rewrite Strategies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define a possible rewrite rule for a single expression, and then have that execute repeatedly on the graph. \n",
    "\n",
    "We can also combine several strategies and execute them together:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import metadsl_rewrite\n",
    "\n",
    "@metadsl_rewrite.rule\n",
    "def _add(x: int, y: int):\n",
    "    return create_object(x) + create_object(y), lambda: create_object(x + y)\n",
    "\n",
    "@metadsl_rewrite.rule\n",
    "def _do_things(x: int):\n",
    "    return create_object(x).do_things(), lambda: create_object(x * 2)\n",
    "\n",
    "\n",
    "inner_strategy = metadsl_rewrite.StrategySequence(\n",
    "    _add,\n",
    "    _do_things\n",
    ")\n",
    "strategies = metadsl_rewrite.StrategyRepeat(\n",
    "    metadsl_rewrite.StrategyFold(\n",
    "        inner_strategy\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The requirements for these replacements is that they take in some arguments, which can match any expression in the graph.\n",
    "Their first return value builds up a template expression based on the inputs, that shows what it should match again. The second\n",
    "is a is the resulting to replace it with. Note that both should have the same type, because if you replace an expression it should not invalidate\n",
    "the type of something it is a part of:\n",
    "\n",
    "We can call these on an expression and it will return a replaced version of it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyObject.__add__(create_object(123), create_object(123))\n",
      "create_object(246)\n"
     ]
    }
   ],
   "source": [
    "print(o+o)\n",
    "print(metadsl_rewrite.execute(o + o, strategies))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also have a rule that unwraps the `int` from the object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "246"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@metadsl.expression\n",
    "def unwrap_object(o: MyObject) -> int:\n",
    "    ...\n",
    "    \n",
    "@metadsl_rewrite.rule\n",
    "def _unwrap_object(i: int):\n",
    "    return unwrap_object(create_object(i)), i\n",
    "\n",
    "inner_strategy.strategies = inner_strategy.strategies + (_unwrap_object,)\n",
    "\n",
    "metadsl_rewrite.execute(unwrap_object(o + o), strategies)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is nice, because now we can write our unboxing as a replacement, which means it's nice and type safe."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.0 ('metadsl')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "8f1320fcc4eb9d472fe42e82e7e653b438a14df950aaaa73e390cd48905e514f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
